package cn.itcast.tags.ml.classification

import org.apache.spark.ml.classification.{LogisticRegression, LogisticRegressionModel}
import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator
import org.apache.spark.ml.feature.{OneHotEncoder, OneHotEncoderEstimator, StringIndexer, VectorAssembler}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}
import org.apache.spark.sql.functions._
object TitanicLrClassFication {

  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession.builder()
      .appName(this.getClass.getSimpleName.stripSuffix("$"))
      .master("local[4]")
      .config("spark.sql.shuffle.partitions", 4)
      .getOrCreate()

    val rawTitanicDF: DataFrame = spark.read
      .option("header", true)
      .option("inferSchema", true)
      .csv("datas/titanic/train.csv")

    import spark.implicits._
    //TODO: 2.数据准备，特征工程（提取，转换与选择)
    val avgAge: Double = rawTitanicDF
      .select($"Age")
      .filter($"Age".isNotNull)
      .select(
        round(avg($"Age"), 2).as("avgAge")
      )
      .first()
      .getAs[Double](0)


    val ageTitanicDF: DataFrame = rawTitanicDF
      .select(
        $"Survived".as("label"),
        $"Pclass", $"Sex", $"SibSp", $"Parch", $"Fare", $"Age",
        //当为null时替换
        when($"Age".isNotNull, $"Age").otherwise(avgAge).as("defaultAge")
      )

    //对Sex字段类别进行转换
    //male->0   female->1
    val indexer: StringIndexer = new StringIndexer()
      .setInputCol("Sex")
      .setOutputCol("sexIndex")

    val indexerTitanicDF: DataFrame = indexer.fit(ageTitanicDF).transform(ageTitanicDF)
    //male->[1.0,0.0] female->[0.0,1.0]
    val encoder: OneHotEncoder = new OneHotEncoder()
      .setInputCol("sexIndex")
      .setOutputCol("sexVector")
      .setDropLast(false)

    val sexTitanicDF: DataFrame = encoder.transform(indexerTitanicDF)

    //将特征值组合，使用VectorAssembler
    val assembler: VectorAssembler = new VectorAssembler()
      .setInputCols(
        Array("Pclass", "sexVector", "SibSp", "Parch", "Fare", "defaultAge")
      )
      .setOutputCol("features")

    val titanicDF: DataFrame = assembler.transform(sexTitanicDF)

    titanicDF.show(10,false)

    //2.4划分数据集为训练集测试集
    val Array(trainingDF,testingDF) = titanicDF.randomSplit(Array(0.8, 0.2))
    trainingDF.cache().count()

    trainingDF.show(10,false)

    //使用算法和数据构建模型
    val logisticRegression: LogisticRegression = new LogisticRegression()
      .setLabelCol("label")
      .setFeaturesCol("features")
      .setPredictionCol("prediction")//使用模型预测时预测值的列名称
      .setFamily("binomial") //二分类
      .setStandardization(true)
      //超参数
      .setMaxIter(100)
      .setRegParam(0.1)
      .setElasticNetParam(0.8)

    val lrModel: LogisticRegressionModel = logisticRegression.fit(trainingDF)
    println(lrModel.coefficientMatrix) //斜率
    println(lrModel.interceptVector)  //截距


    //模型评估
    val predictiondf: DataFrame = lrModel.transform(testingDF)
    predictiondf
      .select("label","prediction","probability","features")
      .show(40,false)

    val accuracy: Double = new MulticlassClassificationEvaluator()//多级分类评估者
      .setLabelCol("label")
      .setPredictionCol("prediction")
      .setMetricName("accuracy")
      .evaluate(predictiondf)

    println(s"准确率:${accuracy}")

    spark.stop()

  }

}
